library(sandwich)
library(lmtest)
#library(parallel)
library(foreach)
library(doSNOW)

# From Aronow, Samii, and Assenova (2015)
# Cluster robust variance estimation function
robust.se.nodfc <- function(model, cluster, efres){
  require(sandwich)
  require(lmtest)
  M <- length(unique(cluster))
  N <- length(cluster)
  K <- model$rank
  dfc <- 1
  if (-99 %in% cluster){
    uj <- matrix(NA, nrow = M, ncol = ncol(efres))
    if (sum(cluster == -99) > 1){
      uj[1,] <- apply(efres[which(cluster == -99),], 2, sum)
    } else {
      uj[1,] <- efres[which(cluster == -99),]}
    uj[-1,] <- efres[which(cluster != -99),]
  } else{uj <- apply(efres, 2, function(x) tapply(x, cluster, sum))}
  rcse.cov <- dfc * sandwich(model, meat = crossprod(uj)/N)
  return(rcse.cov)
}

# Dyadic cluster robust via multiway decomposition
dyad.robust.se <- function(dataUp, fit, index, dyad.mat){
  efres <- estfun(fit)
  
  getdcrUp <- function(i){
    clusUp <-apply(dyad.mat, 1, function(x) as.numeric (index[i] %in% x))
    clusIndexUp <- clusUp*-99 + (1-clusUp)*(1:nrow(dyad.mat))
    return(robust.se.nodfc(fit, clusIndexUp, efres))  
  }
  print("started")
  dcrUp <- foreach(i = 1:length(index), .combine = '+', 
                   .export = c("robust.se.nodfc"),
                   .packages = c("sandwich", "lmtest")) %dopar% {
                     sink(file = "cueingprogress.txt", append = TRUE)
                     print(i)
                     sink()
                     getdcrUp(i)
                   }
  
  # substract naive CR:
  dcrUp2 <- dcrUp - robust.se.nodfc(fit, dataUp$dyads, efres)
  # substract HR:
  Vhat <- dcrUp2 - (length(index)-2)*vcovHC(fit,type="HC0")
  return(Vhat)
}